# 用于文档材料的总结提炼
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate,SystemMessagePromptTemplate,HumanMessagePromptTemplate
# from utils.log_manager import app_logger
from PyPDF2 import PdfReader
from langchain.text_splitter import CharacterTextSplitter

class AbstractModel:
    def __init__(self,api_key,api_base,model_name="glm-4-flash"):
        self.llm = ChatOpenAI(
            temperature=0.95,
            model=model_name,
            openai_api_key=api_key,
            openai_api_base=api_base
        )
        self.set_prompt()
    
    def set_prompt(self):
        self.system_template = """
        从以下文档中提炼出关键信息，简洁地总结内容，但确保保留核心要点。
        请避免冗长的描述，并突出关键信息，以便后续处理生成试题。
        请以简明扼要的形式呈现内容，重点关注以下方面：
        1. 主题
        2. 重要概念
        3. 主要论点或结论
        4. 必要的定义、公式或数据
        5. 适合用作试题的核心内容

        请确保输出内容清晰、易懂
        """
        self.prompt = ChatPromptTemplate(
            messages=[
                SystemMessagePromptTemplate.from_template(
                    self.system_template,
                ),
                # 缓存占位符
                # MessagesPlaceholder(variable_name="chat_history"),
                HumanMessagePromptTemplate.from_template("文档内容：{content}")
            ]
        )
        self.conversation = self.prompt | self.llm
    
    def generate_abstract(self,content):
        """生成摘要"""
       
        result = self.conversation.invoke({"content":content})
        return result.content
    

if __name__ == "__main__":
    api_key = "sk-51dd4f350cdd49f58083396702d29b25"
    api_base = "https://api.deepseek.com"
    model = AbstractModel(api_key,api_base)
    model.set_prompt()

    # 读取PDF文件的第30页内容
    pdf_path = "E:\\APP\\QuickLearn\\server\\static\\docs\\CISSP官方学习手册（第9版） ([美]迈克·查普尔(Mike Chapple)) (Z-Library).pdf"
    reader = PdfReader(pdf_path)
    content = ""
    for i in range(1,10):
        page = reader.pages[i]  # 第30页的索引为29
        content += page.extract_text()

    # 使用CharacterTextSplitter将内容分割成小段
    text_splitter = CharacterTextSplitter(chunk_size=2000, chunk_overlap=0)
    texts = text_splitter.split_text(content)

    # 统计token数量
    token_count = model.llm.get_num_tokens(texts[0])
    print(f"传递给模型的Token数量: {token_count}")

    abstract = model.generate_abstract(content)
    print(abstract)
